import os
import sys
from multiprocessing import cpu_count
import time
from itertools import islice

from gensim.corpora.dictionary import Dictionary
from gensim.models.ldamulticore import LdaMulticore
from gensim.corpora import TextCorpus
from gensim.models.wrappers.ldamallet import LdaMallet, malletmodel2ldamodel

current_folder = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.dirname(os.path.dirname(current_folder)))

from news_classification import mallet_bin, USE_MALLET
from news_classification.utils import filter_words


def doc_generator1(corpus_fp, do_filter=False):
    start = time.perf_counter()
    counter = 0
    with open(corpus_fp, encoding='utf-8') as fi:
        for line in fi:
            words = [w.strip() for w in line.strip().split()]
            if do_filter:
                words = filter_words(words)
            if len(words) == 0:
                continue
            counter += 1
            if counter % 1000 == 0:
                print('\r%d/%fs..' % (counter, time.perf_counter() - start), end='')
            yield words


def doc_generator2(corpus_fp, do_filter=False):
    start = time.perf_counter()
    counter = 0
    with open(corpus_fp, encoding='utf-8') as fi:
        while 1:
            lines = list(islice(fi, 10000))
            if not lines:
                break
            for line in lines:
                words = [w.strip() for w in line.strip().split()]
                if do_filter:
                    words = filter_words(words)
                if len(words) == 0:
                    continue
                counter += 1
                if counter % 1000 == 0:
                    print('\r%d/%fs..' % (counter, time.perf_counter() - start), end='')
                yield words


doc_generator = doc_generator2


def train(corpus_fp, save_model_fp, topic_num=100):
    """
    训练主题模型，挖掘新闻中的主题，辅助分类。
    """
    print('Generating dictionary..')
    dictionary = Dictionary(doc_generator(corpus_fp), 200000)
    print('Training lda model..')
    corpus = TextCorpus(corpus_fp, dictionary=dictionary, token_filters=[])
    if USE_MALLET:
        print('Using Mallet LDA..')
        mallet_lda = LdaMallet(mallet_bin, corpus=corpus, num_topics=topic_num, id2word=dictionary,
                               workers=cpu_count() - 1, iterations=50)
        print('Converting Mallet LDA to gensim LDA..')
        lda = malletmodel2ldamodel(mallet_lda)
    else:
        lda = LdaMulticore(corpus, num_topics=topic_num, id2word=dictionary, workers=int(2 * cpu_count() / 3))
    print('Saving lda model..')
    lda.save(save_model_fp)


if __name__ == '__main__':
    train(sys.argv[1], sys.argv[2], int(sys.argv[3]))
